#!/usr/bin/env bash

# Read Slurm user accounts with sacctmgr and create a user_settings.conf config file
# to be copied to the /etc/slurm/ directory.
# Homepage: https://github.com/OleHolmNielsen/Slurm_tools/

# Syntax of the user_settings.conf config file is:
# [DEFAULT|UNIX_group|username]:[Type]:value
# Type may be: fairshare, GrpTRES, GrpTRESRunMins etc.

# Run commands remotely if your Slurm host does not have gawk version 4:
# Execute commands on a remote host by password-less SSH:
# export remote="ssh <slurm-host>"

# Slurm commands:
export sacctmgr=/usr/bin/sacctmgr

# WARNING: GNU gawk version 4.0 or later is required for arrays of arrays
awk_version=`awk --version | head -1 | awk '{version=$3; split(version,v,"."); print v[1]}'`
if test "$awk_version" = "3"
then
	echo -n "Sorry, gawk version 4.0 or later is required.  Your version is: "
	awk --version | head -1
	exit 1
fi

# Header of sacctmgr -snorp show associations:
# Cluster|Account|User|Partition|Share|GrpJobs|GrpTRES|GrpSubmit|GrpWall|GrpTRESMins|MaxJobs|MaxTRES|MaxTRESPerNode|MaxSubmitJobs|MaxWall|MaxTRESMins|QOS|Def QOS|GrpTRESRunMins|

# Get the Slurm version
SLURM_VERSION="`$sacctmgr --version | awk '{print $2}'`"
# Version comparison, https://apple.stackexchange.com/questions/83939/compare-multi-digit-version-numbers-in-bash/
function version { echo "$@" | awk -F. '{ printf("%03d%03d%03d", $1,$2,$3); }'; }
 
# sacctmgr show associations format (non-default) adding more recent Slurm parameters
# Default format: Cluster|Account|User|Partition|Share|GrpJobs|GrpTRES|GrpSubmit|GrpWall|GrpTRESMins|MaxJobs|MaxTRES|MaxTRESPerNode|MaxSubmitJobs|MaxWall|MaxTRESMins|QOS|Def QOS|GrpTRESRunMins|
# See "man acctmgr" section "LIST/SHOW ASSOCIATION FORMAT OPTIONS"
export assocfmt="Cluster,Account,User,Partition,Fairshare,GrpJobs,GrpTRES,GrpSubmit,GrpWall,GrpTRESMins,MaxJobs,MaxTRES,MaxTRESPerNode,MaxSubmitJobs,MaxWall,MaxTRESMins,QOS,DefaultQOS,GrpTRESRunMins"
# Add parameters from newer Slurm versions
if test $(version $SLURM_VERSION) -ge $(version 18.08)
then
        assocfmt="$assocfmt,GrpJobsAccrue,MaxJobsAccrue"
        echo "### Slurm $SLURM_VERSION Format=$assocfmt"
fi

$remote $sacctmgr -snorp show associations format=$assocfmt | awk '
BEGIN {
	debug = 0	# Debugging flag
	# IGNORECASE = 1	# Make case-insensitive string comparisons (for TRES comparisons)

	# Create an array of Slurm factors
	assocfmt=ENVIRON["assocfmt"]
	split(assocfmt, slurm_factors, ",")

	# Print a header
	print "###"
	print "### Syntax of the user_settings.conf config file is:"
	print "### [DEFAULT|UNIX_group|username]:[Type]:value"
	print "### Type: " string
	print "###"

	FS="|"	# Set the Field Separator
}
{
	# Read list of existing accounts (parseable output)
	if (debug > 1) print "Got association " $0
	a = $2	# account
	u = $3	# user
	if (a == "root") next	# Skip the root account
	if (u == "") next	# Skip non-user accounts
	accounts[a]	= a
	useraccount[u]	= a
	user[u]		= u

	# Record the user settings
	setting[u]["Fairshare"]		= $5
	setting[u]["GrpJobs"]		= $6
	setting[u]["GrpTRES"]		= tolower($7)	# cpu= must be in lowercase
	setting[u]["GrpTRESMins"]	= $10
	setting[u]["MaxTRES"]		= $12
	setting[u]["MaxTRESPerNode"]	= $13
	setting[u]["MaxTRESMins"]	= $16
	setting[u]["QOS"]		= $17
	setting[u]["DefaultQOS"]	= $18
	setting[u]["GrpTRESRunMins"]	= tolower($19)	# cpu= must be in lowercase
	if (NF >= 20) {		# New factors from Slurm 18.08
		setting[u]["GrpJobsAccrue"]	= $20
		setting[u]["MaxJobsAccrue"]	= $21
	}
	for (i in slurm_factors) {
		sf = slurm_factors[i]
		# if (setting[u][sf] != "") {
		# Allow for empty string values:
		{
			# Count the setting values
			v = setting[u][sf]
			if (debug > 0) printf("%s:%s:%s\n", u, sf, v)
			# Accumulate account settings
			counts[a][sf][v]++
			value[a][sf][v] = v
			# Accumulate global settings
			globalcounts[sf][v]++
			globalvalue[sf][v] = v
		}
	}
} END {
	print "###"
	print "### The global DEFAULT values of settings"
	for (j in slurm_factors) {
		sf = slurm_factors[j]
		# Special case: The QOS default is always the "normal" QOS
		if (sf == "QOS" || sf == "DefaultQOS") {
			globaldefaultvalue[sf] = "normal"
		}
		if (isarray(globalcounts[sf])) {
			# Locate the highest count, assumed to be the default for this account
			cmax = 0
			for (k in globalcounts[sf]) {
				c = globalcounts[sf][k]
				v = globalvalue[sf][k]
				print "### " sf " value=\"" v "\" has count=" c
				if (c > cmax) {
					cmax = c
					globaldefaultvalue[sf] = v
				}
			}
		}
		if (globaldefaultvalue[sf] != "")
			printf("%s:%s:%s\n", "DEFAULT", sf, globaldefaultvalue[sf])
		else
			printf("### Empty default value: %s:%s:%s\n", "DEFAULT", sf, globaldefaultvalue[sf])
	}
	print "###"
	print "### The account default values of settings"
	for (i in accounts) {
		a = accounts[i]
		for (j in slurm_factors) {
			sf = slurm_factors[j]
			defaultvalue[a][sf] = globaldefaultvalue[sf]
			# Special case: The QOS default is always the "normal" QOS
			if (sf == "QOS" || sf == "DefaultQOS") {
				defaultvalue[a][sf] = "normal"
				continue
			}
			if (!isarray(counts[a][sf])) continue
			cmax = 0
			for (k in counts[a][sf]) {
				c = counts[a][sf][k]
				v = value[a][sf][k]
				if (debug > 0) print "# Account " a " setting " sf " value=\"" v "\" has count=" c
				# Locate the highest count, assumed to be the default for this account
				if (c > cmax) {
					cmax = c
					defaultvalue[a][sf] = v
				}
			}
			# Print the value if it differs from the global default
			if (defaultvalue[a][sf] != globaldefaultvalue[sf])
				printf("%s:%s:%s\n", a, sf, defaultvalue[a][sf])
		}
	}
	print "###"
	print "### The user settings, omitting group DEFAULT values"
	for (i in user) {
		u = user[i]
		a = useraccount[u]
		if (debug > 0) print "# User " u " in group " a
		for (j in slurm_factors) {
			sf = slurm_factors[j]
			# Print the value only if it differs from the account default
			if (setting[u][sf] != "" && setting[u][sf] != defaultvalue[a][sf])
				printf("%s:%s:%s\n", u, sf, setting[u][sf])
		}
	}
}'
